import cx_Oracle
import json
from sqlalchemy import create_engine, text, MetaData, Table, Column, String, inspect
from sqlalchemy.types import VARCHAR, DATE, Integer, CLOB
from langchain_community.utilities.sql_database import SQLDatabase
from langchain_core.prompts import FewShotPromptTemplate
from boat_nl2sql.utiles import extract_sql_from_markdown,load_table_info_from_json

# # Oracle 连接配置
# username = "sys"
# password = "123456"
# host = "10.1.161.137"
# port = "1521"
# sid = "ORCL"
# schemas = ["BAYONET_BASICS", "BAYONET_BUSSINESS", "BAYONET_DYNAMIC"]

# Oracle 连接配置
username = "sys"
password = "oracle"
host = "10.1.161.53"
port = "11521"
sid = "helowin"
schemas = ["BAYONET_BASICS", "BAYONET_BUSSINESS", "BAYONET_DYNAMIC"]


def create_oracle_connection():
    """创建 Oracle 数据库连接"""
    dsn = cx_Oracle.makedsn(host, port, sid=sid)
    return cx_Oracle.connect(
        user=username,
        password=password,
        dsn=dsn,
        mode=cx_Oracle.SYSDBA
    )


def get_table_details():
    """直接从 Oracle 数据库获取表的详细信息"""
    try:
        conn = create_oracle_connection()
        cursor = conn.cursor()

        tables_details = {}

        for schema in schemas:
            # 获取该 schema 下的所有表
            cursor.execute(f"""  
                SELECT TABLE_NAME   
                FROM ALL_TABLES   
                WHERE OWNER = '{schema}'  
            """)

            tables = [row[0] for row in cursor.fetchall()]

            for table in tables:
                qualified_table_name = f"{schema}.{table}"

                # 获取列信息
                cursor.execute(f"""  
                    SELECT COLUMN_NAME, DATA_TYPE,   
                           CASE WHEN DATA_TYPE LIKE '%CHAR%' THEN DATA_LENGTH ELSE NULL END AS DATA_LENGTH,  
                           NULLABLE  
                    FROM ALL_TAB_COLUMNS  
                    WHERE OWNER = '{schema}' AND TABLE_NAME = '{table}'  
                    ORDER BY COLUMN_ID  
                """)

                columns = []
                for col_name, data_type, data_length, nullable in cursor.fetchall():
                    # 格式化数据类型
                    if data_length is not None:
                        type_str = f"{data_type}({data_length})"
                    else:
                        type_str = data_type

                        # 添加 NOT NULL 约束
                    nullable_str = " NOT NULL" if nullable == 'N' else ""

                    columns.append(f"{col_name} {type_str}{nullable_str}")

                    # 获取主键信息
                cursor.execute(f"""  
                    SELECT cols.column_name  
                    FROM all_constraints cons, all_cons_columns cols  
                    WHERE cons.constraint_type = 'P'  
                    AND cons.constraint_name = cols.constraint_name  
                    AND cons.owner = '{schema}'  
                    AND cols.owner = '{schema}'  
                    AND cons.table_name = '{table}'  
                    AND cols.table_name = '{table}'  
                """)

                pk_columns = [row[0] for row in cursor.fetchall()]

                # 添加主键标记
                for i, col_def in enumerate(columns):
                    col_name = col_def.split()[0]
                    if col_name in pk_columns:
                        # 替换列定义，添加 PRIMARY KEY
                        columns[i] = col_def + " PRIMARY KEY"

                        # 创建表定义
                create_statement = f"CREATE TABLE {qualified_table_name} (\n  "
                create_statement += ",\n  ".join(columns)
                create_statement += "\n);"

                # 存储表定义
                tables_details[qualified_table_name] = create_statement

                print(f"获取到表 {qualified_table_name} 的结构")

        cursor.close()
        conn.close()

        return tables_details

    except Exception as e:
        print(f"获取表详细信息时出错: {e}")
        import traceback
        traceback.print_exc()
        return {}


def save_table_details_to_json(file_path="oracle_table_info.json"):
    """将表详细信息保存到 JSON 文件"""
    try:
        table_details = get_table_details()

        if not table_details:
            print("未获取到任何表信息，无法保存")
            return False

        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(table_details, f, indent=2)

        print(f"已将 {len(table_details)} 个表的详细信息保存到 {file_path}")
        return True

    except Exception as e:
        print(f"保存表信息到 JSON 文件时出错: {e}")
        return False



def get_tables_from_schemas():
    """获取指定 schema 中所有表的完整名称列表"""
    try:
        conn = create_oracle_connection()
        cursor = conn.cursor()

        all_tables = []
        for schema in schemas:
            cursor.execute(f"""  
                SELECT '{schema}.' || TABLE_NAME   
                FROM ALL_TABLES   
                WHERE OWNER = '{schema}'  
            """)

            schema_tables = [row[0] for row in cursor.fetchall()]
            all_tables.extend(schema_tables)

            if schema_tables:
                print(f"在 {schema} schema 中找到 {len(schema_tables)} 个表")
            else:
                print(f"在 {schema} schema 中没有找到表")

        cursor.close()
        conn.close()

        print(f"总共找到 {len(all_tables)} 个表")
        return all_tables

    except Exception as e:
        print(f"获取表列表时出错: {e}")
        import traceback
        traceback.print_exc()
        return []


class NoReflectMetaData(MetaData):
    """禁用表反射的 MetaData 类"""

    def reflect(self, *args, **kwargs):
        """覆盖反射方法，什么都不做"""
        pass


class OracleSQLDatabase(SQLDatabase):
    """
    自定义 SQLDatabase 类，完全禁用表反射
    而是使用自定义的表信息
    """

    def __init__(
            self,
            engine,
            schema=None,
            metadata=None,
            sample_rows_in_table_info=3,
            indexes_in_table_info=False,
            custom_table_info=None,
            view_support=False,
            include_tables=None,
            exclude_tables=None,
            ignore_tables=None,
    ):
        """
        完全重写初始化方法，避免父类的表反射
        """
        self._engine = engine
        self._schema = schema
        self._include_tables = include_tables
        self._indexes_in_table_info = indexes_in_table_info
        self._sample_rows_in_table_info = sample_rows_in_table_info
        self._view_support = view_support

        # 使用自定义的不反射 MetaData
        self._metadata = NoReflectMetaData()

        # 保存自定义表信息
        self._custom_table_info = custom_table_info

        # 从表定义创建表信息字符串
        self._custom_table_info_str = "\n\n".join(
            [info for info in custom_table_info.values()]
        ) if custom_table_info else ""

        # 表信息缓存
        self._table_info = None

        # 占位符风格
        self._dialect = self._engine.dialect.name

    @property
    def dialect(self):
        """获取数据库方言"""
        return self._dialect

    def _get_all_table_names(self) -> list:
        """获取所有表名"""
        return list(self._custom_table_info.keys()) if self._custom_table_info else []

    def get_usable_table_names(self) -> list:
        """获取可用表名"""
        return self._get_all_table_names()

    @property
    def table_info(self) -> str:
        """获取表信息字符串"""
        if not self._table_info:
            self._table_info = self._custom_table_info_str
        return self._table_info

    def run(self, command: str, fetch: str = "all") -> str:
        """执行 SQL 命令"""
        with self._engine.connect() as connection:
            if self._dialect == "sqlite":
                connection = connection.execution_options(
                    sqlite_raw_colnames=True
                )
            cursor = connection.execute(text(command))
            if cursor.returns_rows:
                if fetch == "all":
                    result = cursor.fetchall()
                elif fetch == "one":
                    result = cursor.fetchone()
                else:
                    raise ValueError("Fetch parameter must be either 'one' or 'all'")
                return result
            return []


def get_db_connection(sample_rows=2, custom_table_info=None, include_tables=None):
    """
    获取 LangChain SQLDatabase 连接，使用自定义表信息

    参数:
        sample_rows (int): 每个表中示例的行数
        custom_table_info (dict): 自定义表信息字典
        include_tables (list): 要包含的表列表，格式为 ['SCHEMA.TABLE']

    返回:
        SQLDatabase 对象
    """
    try:
        # 创建 SQLAlchemy 引擎
        engine = create_engine("oracle://", creator=create_oracle_connection)

        # 如果没有提供自定义表信息，尝试从 JSON 文件加载
        if custom_table_info is None:
            custom_table_info = load_table_info_from_json("oracle_table_info.json")

            # 确保我们有表信息
        if not custom_table_info:
            print("警告: 未加载到任何表信息，尝试从数据库获取")
            custom_table_info = get_table_details()

            if not custom_table_info:
                raise ValueError("无法获取表信息，无法创建 SQLDatabase")

                # 如果指定了 include_tables，只保留这些表的信息
        if include_tables:
            filtered_table_info = {}
            for table in include_tables:
                if table in custom_table_info:
                    filtered_table_info[table] = custom_table_info[table]
                else:
                    print(f"警告: 表 {table} 在自定义表信息中不存在")

            if not filtered_table_info:
                print("错误: 指定的表都不在自定义表信息中")
                print(f"指定的表: {include_tables}")
                print(f"可用的表: {list(custom_table_info.keys())}")
                raise ValueError("无法找到指定的表")

            custom_table_info = filtered_table_info
            print(f"已筛选出 {len(custom_table_info)} 个指定的表")

            # 创建自定义的 SQLDatabase 对象
        db = OracleSQLDatabase(
            engine=engine,
            sample_rows_in_table_info=sample_rows,
            custom_table_info=custom_table_info
        )

        # 确认可用表
        usable_tables = db.get_usable_table_names()
        print(f"可用表数量: {len(usable_tables)}")
        if usable_tables:
            print("可用表:", usable_tables)

        return db

    except Exception as e:
        print(f"创建数据库连接时出错: {e}")
        import traceback
        traceback.print_exc()
        raise


def execute_query(sql):
    """直接执行 SQL 查询并返回结果"""
    try:
        conn = create_oracle_connection()
        cursor = conn.cursor()
        cursor.execute(sql)

        # 获取列名
        columns = [col[0] for col in cursor.description]

        # 获取结果
        rows = cursor.fetchall()

        # 格式化结果
        results = []
        for row in rows:
            result_dict = {}
            for i, col in enumerate(columns):
                result_dict[col] = row[i]
            results.append(result_dict)

        cursor.close()
        conn.close()

        return results

    except Exception as e:
        print(f"执行查询时出错: {e}")
        raise


def test_include_tables():
    """测试指定表列表功能"""
    try:
        print("===== 测试指定表列表 =====")

        # 加载全部表信息
        full_table_info = load_table_info_from_json("oracle_table_info.json")
        print(f"全部表信息中包含 {len(full_table_info)} 个表")

        # 指定要包含的表
        include_tables = [
            'BAYONET_BUSSINESS.BOAT_WARNING',
            'BAYONET_BUSSINESS.OFF_SITE_CASE',
            'BAYONET_BUSSINESS.ZD_WARNING_TYPE',
            'BAYONET_DYNAMIC.VIDEO_ANALYSIS',
            'BAYONET_DYNAMIC.DATAFUSION',
            'BAYONET_BASICS.SYS_BAYONET'
        ]

        # 获取数据库连接
        db = get_db_connection(
            sample_rows=2,
            custom_table_info=full_table_info,
            include_tables=include_tables
        )

        # 验证结果
        usable_tables = db.get_usable_table_names()
        print(f"可用表: {usable_tables}")

        # 检查是否所有指定的表都包含在可用表中
        missing_tables = [t for t in include_tables if t not in usable_tables]
        if missing_tables:
            print(f"警告: 以下表未包含在可用表中: {missing_tables}")
        else:
            print("所有指定的表都已包含在可用表中")

        # 测试查询
        # print("\n===== 测试表查询 =====")
        # for table in usable_tables:
        #     try:
        #         query = f"SELECT * FROM {table} WHERE ROWNUM <= 1"
        #         print(f"\n执行查询: {query}")
        #         result = db.run(query)
        #         if result:
        #             print(f"查询结果: {result}")
        #         else:
        #             print("查询结果为空")
        #     except Exception as e:
        #         print(f"查询表 {table} 时出错: {e}")

        return True

    except Exception as e:
        print(f"测试指定表列表时出错: {e}")
        import traceback
        traceback.print_exc()
        return False


def fix_table_names(query, all_table_names):
    """修复 SQL 查询中的表名，确保它们有正确的 schema 前缀"""
    # 此处添加您原有的表名修复代码
    return query


def test_sql_chain():
    """测试 LangChain 的 SQL 查询链"""
    try:
        from langchain.chains import create_sql_query_chain
        from langchain.prompts import PromptTemplate

        # 在这里导入您的 LLM
        from llm_module import SiliconFlow
        llm = SiliconFlow()

        print("\n===== 测试 SQL 查询链 =====")
        # print("注意: 这个测试需要您提供一个 LLM 实例")
        # print("请在 test_sql_chain 函数中取消注释导入您的 LLM 的代码")
        # 加载全部表信息
        full_table_info = load_table_info_from_json("oracle_table_info.json")
        print(f"全部表信息中包含 {len(full_table_info)} 个表")

        # 指定要包含的表
        include_tables = [
            'BAYONET_BUSSINESS.BOAT_WARNING',
            'BAYONET_BUSSINESS.OFF_SITE_CASE',
            'BAYONET_BUSSINESS.ZD_WARNING_TYPE',
            'BAYONET_DYNAMIC.VIDEO_ANALYSIS',
            'BAYONET_DYNAMIC.DATAFUSION',
            'BAYONET_BASICS.SYS_BAYONET'
        ]

        # 获取数据库连接
        db = get_db_connection(
            sample_rows=2,
            custom_table_info=full_table_info,
            include_tables=include_tables
        )

        # 配置提示模板
        SQL_PROMPT = PromptTemplate.from_template(
            """你是一个SQL专家。根据以下表结构：  
            {table_info}  

            问题：{input}  
            请严格按以下规则响应：  
            1. 只输出SQL代码  
            2. 使用标准Oracle SQL语法  
            3. 不要包含分号结尾  
            4. 最多返回{top_k}条记录  
            5. Oracle中使用ROWNUM来限制返回记录数  

            SQL查询："""
        )
        examples = [
            {
                "input": "查询名称中包含'船闸'的所有卡口基本信息，按名称排序。",
                "query": "SELECT ID, NAME, UNIT_ID, DES, PHOTOHTTP  FROM BAYONET_BASICS.SYS_BAYONET  WHERE NAME LIKE '%船闸%'  ORDER BY NAME",
            },
            {
                "input": "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。",
                "query": "SELECT * FROM (SELECT b.NAME AS 卡口名称, COUNT(d.CODE) AS 通过船舶数量  FROM BAYONET_BASICS.SYS_BAYONET b  JOIN BAYONET_DYNAMIC.DATAFUSION d ON b.ID = d.BAYONET_ID  WHERE d.PASSTIME >= TRUNC(SYSDATE, 'YYYY') AND d.PASSTIME < ADD_MONTHS(TRUNC(SYSDATE, 'YYYY'), 1) GROUP BY b.NAME  ORDER BY 通过船舶数量 DESC) WHERE ROWNUM <= 10",
            },
            {
                "input": "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。",
                "query": "WITH warning_counts AS (  SELECT t.NAME AS 预警类型,   COUNT(w.CODE) AS 预警次数  FROM BAYONET_BUSSINESS.ZD_WARNING_TYPE t  JOIN BAYONET_BUSSINESS.BOAT_WARNING w ON t.ID = w.WARNING_TYPE  GROUP BY t.NAME  )  SELECT    预警类型,预警次数, ROUND(预警次数 * 100 / SUM(预警次数) OVER(), 2) || '%' AS 占比 FROM warning_counts ORDER BY 预警次数 DESC "
            },
            {
                "input": "查询名为'苏无锡货08168'的船舶最近一次通过系统的完整记录，包括通过卡口信息、预警情况和消息发送记录。",
                "query": "SELECT * FROM (SELECT d.PASSTIME AS 通过时间, b.NAME AS 卡口名称, d.BOAT_NAME AS 船舶名称, wt.NAME AS 预警类型, w.WARNING_TIME AS 预警时间, w.WARNING_STATUS AS 预警状态,  COUNT(m.ID) AS 发送消息数量  FROM BAYONET_DYNAMIC.DATAFUSION d  JOIN BAYONET_BASICS.SYS_BAYONET b ON d.BAYONET_ID = b.ID  LEFT JOIN BAYONET_BUSSINESS.BOAT_WARNING w ON d.CODE = w.EVENTCODE  LEFT JOIN BAYONET_BUSSINESS.ZD_WARNING_TYPE wt ON w.WARNING_TYPE = wt.ID  LEFT JOIN BAYONET_BUSSINESS.MESSAGE_SEND_RECORD m ON w.WARNCODE = m.WARNING_CODE  WHERE d.BOAT_NAME = '苏无锡货08168'  GROUP BY d.PASSTIME, b.NAME, d.BOAT_NAME, wt.NAME, w.WARNING_TIME, w.WARNING_STATUS  ORDER BY d.PASSTIME DESC) WHERE ROWNUM = 1"
            }

        ]

        few_shot_example_prompt = PromptTemplate.from_template("输入问题: {input}\n生成SQL: {query}")
        few_shot_prompt = FewShotPromptTemplate(
            examples=examples,
            example_prompt=few_shot_example_prompt,
            prefix="你是一个Oracle SQL专家。请根据以下表结构生成对应问题的Oracle SQL语句，且尽量符合标准语法。\n\n以下是表结构（table_info）：\n{table_info}\n\n这是一些问题及其对应的SQL样例:",
            suffix="\n---\n问题: {input}\n对应SQL: ",
            input_variables=["input", "table_info"],
        )


        # 创建 SQL 查询链
        generate_query = create_sql_query_chain(llm=llm, db=db, prompt=SQL_PROMPT, k=5)

        # 示例查询
        question = "查询所有卡口的名称"
        print(f"问题: {question}")

        few_shot_sql = few_shot_prompt.format(input=question, table_info=db.table_info)

        llm_response = llm.invoke(few_shot_sql)
        print("FewShot 生成的 SQL:", llm_response)

        if "```sql" in llm_response:
            extracted_sql = extract_sql_from_markdown(llm_response)
            print("提取出的 SQL:")
            print(extracted_sql)
        else:
            extracted_sql = llm_response.replace(";","")

        try:
            # 执行SQL
            sql_result = db.run(extracted_sql)
            print("SQL结果:", sql_result)

            TEXT_PROMPT = PromptTemplate.from_template(
                """根据问题{test_question}，和查询数据库返回的结果{sql_result}，用简介的语言，回答提出的问题。""")
            # 格式化 PromptTemplate 为字符串
            formatted_prompt = TEXT_PROMPT.format(test_question=question, sql_result=sql_result)

            result = llm.siliconflow_completions(model="deepseek-ai/DeepSeek-V3", prompt=formatted_prompt)
            print("结果:", result)

        except Exception as e:
            print(f"执行查询时出错: {e}")

        # 生成 SQL
        # generated_sql = generate_query.invoke({"question": question})
        # print(f"生成的SQL: {generated_sql}")
        # try:
        #     result = db.run(generated_sql)
        #     print(f"查询结果: {result}")
        # except Exception as e:
        #     print(f"执行查询时出错: {e}")


        # 测试表查询
        # test_queries = [
        #     "SELECT * FROM BAYONET_BASICS.SYS_BAYONET WHERE ROWNUM <= 3",
        #     "SELECT * FROM BAYONET_BUSSINESS.BOAT_WARNING WHERE ROWNUM <= 3",
        #     "SELECT * FROM BAYONET_DYNAMIC.DATAFUSION WHERE ROWNUM <= 3"
        # ]
        #
        # for query in test_queries:
        #     print(f"\n执行查询: {query}")
        #     try:
        #         result = db.run(query)
        #         print(f"查询结果: {result}")
        #     except Exception as e:
        #         print(f"执行查询时出错: {e}")

        return True

    except ImportError:
        print("无法导入 LangChain 包，请确保已安装")
        return False
    except Exception as e:
        print(f"测试 SQL 查询链时出错: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_port_api_sql():
    # 加载全部表信息
    full_table_info = load_table_info_from_json("oracle_table_info.json")
    print(f"全部表信息中包含 {len(full_table_info)} 个表")

    # 指定要包含的表
    include_tables = [
        'BAYONET_BUSSINESS.BOAT_WARNING',
        'BAYONET_BUSSINESS.OFF_SITE_CASE',
        'BAYONET_BUSSINESS.ZD_WARNING_TYPE',
        'BAYONET_DYNAMIC.VIDEO_ANALYSIS',
        'BAYONET_DYNAMIC.DATAFUSION',
        'BAYONET_BASICS.SYS_BAYONET'
    ]

    # 获取数据库连接
    db = get_db_connection(
        sample_rows=2,
        custom_table_info=full_table_info,
        include_tables=include_tables
    )

    # 验证结果
    usable_tables = db.get_usable_table_names()
    print(f"可用表: {usable_tables}")


    query = """select sys.name
    bayonetName,
    count(1) total,
    count(case when info.DIRECTION = 1 then '出沪' end) outCount,
    count(case when info.DIRECTION = 2 then '入沪' end) inCount,
    count(case when mvw.SHIP_TYPE_NAME = '干货船' then '干货船' end) dryShip,
    count(case when mvw.SHIP_TYPE_NAME = '油船' then '油船' end) oilShip,
    count(case when mvw.SHIP_TYPE_NAME = '散货船' then '散货船' end) bulkShip,
    count(case when mvw.SHIP_TYPE_NAME = '散装化学品船' then '散装化学品船' end)
    bulkChemicalShip,
    count(case when mvw.SHIP_TYPE_NAME = '散装水泥运输船' then '散装水泥运输船' end)
    bulkCementShip,
    count(case when mvw.SHIP_TYPE_NAME = '多用途船' then '多用途船' end) multipurposeShip,
    count(case
    when mvw.SHIP_TYPE_NAME not in
    ('干货船', '油船', '散货船', '散装化学品船', '散装水泥运输船', '多用途船')
    then '其他' end) other,
    sum(mvw.SHIP_GROSSTON) zd
    from BAYONET_DYNAMIC.DATAFUSION info
    left join GHJG_BASICS.MVW_BOAT_INFO mvw on mvw.SHIP_NAME_CN = info.BOAT_NAME
    LEFT JOIN BAYONET_BASICS.SYS_BAYONET sys ON sys.id = info.BAYONET_ID
    where info.BAYONET_ID is not null
    AND   sys.NAME = '油墩港(东大盈船闸)'
    AND info.PASSTIME >= TO_DATE('2025-01-01', 'YYYY-MM-DD')
    AND info.PASSTIME <= TO_DATE('2025-01-03', 'YYYY-MM-DD')
    group by sys.name"""


#     query = """SELECT
#     sb.NAME AS 卡口,
#     df.CODE,
#     df.DIRECTION,
#     osc.BOAT_TYPE
#  FROM
#     BAYONET_BASICS.SYS_BAYONET sb
# LEFT JOIN
#     BAYONET_DYNAMIC.DATAFUSION df ON sb.ID = df.BAYONET_ID
# LEFT JOIN
#     BAYONET_BUSSINESS.OFF_SITE_CASE osc ON osc.BAYONET_ID = sb.ID
# WHERE
#     sb.NAME = '油墩港(东大盈船闸)'
#     AND df.PASSTIME >= TO_DATE('2025-01-01', 'YYYY-MM-DD')
#     AND df.PASSTIME <= TO_DATE('2025-01-03', 'YYYY-MM-DD')
# """

    print(f"\n执行查询: {query}")
    result = db.run(query)
    if result:
        print(f"查询结果: {result}")
        return True
    else:
        print("查询结果为空")
        return False




if __name__ == "__main__":
    print("===== Oracle 数据库连接测试 =====")

    # result1 = test_include_tables()
    # print(f"\n指定表列表测试: {'成功' if result1 else '失败'}")

    # # 取消注释下面的行以测试 SQL 查询链
    # result2 = test_sql_chain()
    # print(f"\nSQL 查询链测: {'成功' if result2 else '失败'}")

    result3 = test_port_api_sql()
    print(f"\nSQL 查询测试: {'成功' if result3 else '失败'}")